package crawler

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/http/cookiejar"
	"strconv"
	"strings"
	"sync"

	"gitee.com/xiaoa7/processbar"
	"golang.org/x/net/proxy"
	"golang.org/x/text/encoding/simplifiedchinese"
	"golang.org/x/text/transform"
)

type (
	//HtppClient 封装的HttpClient
	HtppClient struct {
		Cfg    *HttpClientConfig
		Jar    *cookiejar.Jar
		Client *http.Client
		Bar    *processbar.Bar
	}
	//HttpClientConfig HtppClient对象配置
	HttpClientConfig struct {
		Socks5ProxyAddress string
		ShowProcessBar     bool
		Header             map[string]string
	}
)

//NewHttpClient 创建定制客户端
func NewHttpClient(cfg *HttpClientConfig) (*HtppClient, error) {
	// Cookie
	jar, _ := cookiejar.New(nil)
	// setup a http client
	httpTransport := &http.Transport{
		DisableKeepAlives: true,
	}
	httpClient := &http.Client{
		Transport: httpTransport,
		Jar:       jar,
	}

	if cfg.Socks5ProxyAddress != "" {
		dialer, err := proxy.SOCKS5("tcp", cfg.Socks5ProxyAddress, nil, proxy.Direct)
		if err != nil {
			return nil, err
		}
		// set our socks5 as the dialer
		httpTransport.Dial = dialer.Dial
	}
	return &HtppClient{Cfg: cfg,
		Jar:    jar,
		Client: httpClient,
		Bar: processbar.NewBar(20,
			processbar.Bar_Front_Color_Green,
			processbar.Show_Length_Mode_MByte),
	}, nil
}

//Get Get Method获取网页内容
func (mhc *HtppClient) Get(url, encoding string) (ByteContent, error) {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}
	if mhc.Cfg.Header != nil {
		for k, v := range mhc.Cfg.Header {
			req.Header.Add(k, v)
		}
	}
	if resp, err := mhc.Client.Do(req); err != nil {
		return nil, err
	} else {
		var body ByteContent
		var err error
		if strings.ToLower(encoding) == "gbk" {
			reader := transform.NewReader(resp.Body, simplifiedchinese.GBK.NewDecoder())
			body, err = ioutil.ReadAll(reader)
		} else {
			body, err = ioutil.ReadAll(resp.Body)
		}
		resp.Body.Close()
		return body, err
	}
}

//Len HEAD Method获取ContentLength
func (mhc *HtppClient) Len(url string) (int, error) {
	req, err := http.NewRequest("HEAD", url, nil)
	if err != nil {
		return 0, err
	}
	for k, v := range mhc.Cfg.Header {
		req.Header.Add(k, v)
	}
	if resp, err := mhc.Client.Do(req); err != nil {
		return 0, err
	} else {
		contentLength, _ := strconv.Atoi(resp.Header.Get("Content-Length"))
		resp.Body.Close()
		return contentLength, err
	}
}

//CloneHttpClient 分拆多个httpclient，方便线程下载，继承原client的cookie和代理设置
func (mhc *HtppClient) CloneHttpClient() (*http.Client, error) {
	httpTransport := &http.Transport{
		DisableKeepAlives: true,
	}
	if mhc.Cfg.Socks5ProxyAddress != "" {
		dialer, err := proxy.SOCKS5("tcp", mhc.Cfg.Socks5ProxyAddress, nil, proxy.Direct)
		if err != nil {
			return nil, err
		}
		httpTransport.Dial = dialer.Dial
	}
	return &http.Client{
		Transport: httpTransport,
		Jar:       mhc.Jar,
	}, nil
}

//DownloadWithMulThreads 多线程下载
func (mhc *HtppClient) DownloadWithMultiThreads(url string, /*网页地址*/
	writer io.Writer,
	threads, /*并发线程数*/
	blockCap /*分段大小 单位Mb*/ int) {
	contentLength, _ := mhc.Len(url)
	blockSize := blockCap * 1024 * 1024 //分块大小
	if contentLength == 0 {
		return
	}
	blocks := (contentLength + blockSize - 1) / blockSize
	blockRanges := make([][2]int, 0)
	pos, end := 0, 0
	cache := map[int]ByteContent{}
	concurrentLock := make(chan bool, threads) //并发锁
	mapLock := new(sync.RWMutex)
	wg := new(sync.WaitGroup)
	if mhc.Cfg.ShowProcessBar {
		mhc.Bar.Reset(int64(contentLength), 0)
	}
	for i := 0; i < blocks; i++ {
		end = pos + blockSize
		if end > contentLength {
			end = contentLength
		}
		blockRanges = append(blockRanges, [2]int{pos, end})
		wg.Add(1)
		concurrentLock <- true
		go mhc.downloadWithRange(url, pos, end, i, cache,
			concurrentLock, mapLock, wg)
		pos = end + 1
		if pos >= contentLength {
			break
		}

	}
	wg.Wait()
	if mhc.Cfg.ShowProcessBar {
		mhc.Bar.Finish()
	}
	//
	for i := 0; i < blocks; i++ {
		writer.Write(cache[i])
	}
}

//downloadWithRange 分段下载
func (mhc *HtppClient) downloadWithRange(url string,
	start, end, index int,
	ret map[int]ByteContent,
	concurrentLock chan bool,
	mapLock *sync.RWMutex,
	wg *sync.WaitGroup) {
	defer func() {
		wg.Done()
		<-concurrentLock
	}()
	for retry := 0; retry < 5; retry++ {
		client, _ := mhc.CloneHttpClient()
		req, err := http.NewRequest("GET", url, nil)
		if err != nil {
			continue
		}
		for k, v := range mhc.Cfg.Header {
			req.Header.Add(k, v)
		}
		req.Header.Set("Range", fmt.Sprintf("bytes=%d-%d", start, end))
		if resp, err := client.Do(req); err != nil {
			continue
		} else {
			buf := bytes.NewBuffer(nil)
			var err error
			if mhc.Cfg.ShowProcessBar {
				_, err = io.Copy(io.MultiWriter(buf, mhc.Bar), resp.Body)
			} else {
				_, err = io.Copy(buf, resp.Body)
			}
			if err != nil {
				continue
			}
			resp.Body.Close()
			mapLock.Lock()
			ret[index] = buf.Bytes()
			mapLock.Unlock()
			break
		}
	}
}
